{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(writing-custom-steps)=\n",
    "# Writing custom steps\n",
    "\n",
    "<!-- Add  overview -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Graph executes built-in task classes, or task classes and functions that you implement.\n",
    "The task parameters include the following:\n",
    "* `class_name` (str): the relative or absolute class name.\n",
    "* `handler` (str): the function handler (if class_name is not specified it is the function handler).\n",
    "* `**class_args`: a set of class `__init__` arguments.\n",
    "\n",
    "For example, see the following simple `echo` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlrun"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mlrun: start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# echo class, custom class example\n",
    "class Echo:\n",
    "    def __init__(self, context, name=None, **kw):\n",
    "        self.context = context\n",
    "        self.name = name\n",
    "        self.kw = kw\n",
    "\n",
    "    def do(self, x):\n",
    "        print(\"Echo:\", self.name, x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mlrun: end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test the graph: first convert the code to function, and then add the step to the graph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 2.40.1 (20161225.0304)\n",
       " -->\n",
       "<!-- Title: mlrun&#45;flow Pages: 1 -->\n",
       "<svg width=\"246pt\" height=\"44pt\"\n",
       " viewBox=\"0.00 0.00 245.89 44.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 40)\">\n",
       "<title>mlrun&#45;flow</title>\n",
       "<polygon fill=\"#ffffff\" stroke=\"transparent\" points=\"-4,4 -4,-40 241.8874,-40 241.8874,4 -4,4\"/>\n",
       "<!-- _start -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>_start</title>\n",
       "<polygon fill=\"#d3d3d3\" stroke=\"#000000\" points=\"38.5476,-.0493 40.698,-.1479 42.8263,-.2953 44.9236,-.4913 46.9815,-.7353 48.9917,-1.0266 50.9463,-1.3645 52.8377,-1.7479 54.6587,-2.1759 56.4025,-2.6472 58.0628,-3.1606 59.634,-3.7147 61.1107,-4.308 62.4882,-4.9388 63.7625,-5.6054 64.9302,-6.3059 65.9882,-7.0385 66.9343,-7.8012 67.7669,-8.5918 68.4849,-9.4082 69.0878,-10.2481 69.5758,-11.1093 69.9496,-11.9894 70.2102,-12.886 70.3595,-13.7965 70.3997,-14.7186 70.3334,-15.6497 70.1636,-16.5873 69.8937,-17.5287 69.5276,-18.4713 69.0691,-19.4127 68.5225,-20.3503 67.8923,-21.2814 67.1831,-22.2035 66.3996,-23.114 65.5464,-24.0106 64.6285,-24.8907 63.6504,-25.7519 62.617,-26.5918 61.5329,-27.4082 60.4024,-28.1988 59.2299,-28.9615 58.0197,-29.6941 56.7755,-30.3946 55.5012,-31.0612 54.2002,-31.692 52.8757,-32.2853 51.5309,-32.8394 50.1684,-33.3528 48.7908,-33.8241 47.4003,-34.2521 45.9989,-34.6355 44.5886,-34.9734 43.1708,-35.2647 41.7472,-35.5087 40.3189,-35.7047 38.8872,-35.8521 37.4531,-35.9507 36.0175,-36 34.5815,-36 33.146,-35.9507 31.7119,-35.8521 30.2801,-35.7047 28.8519,-35.5087 27.4282,-35.2647 26.0105,-34.9734 24.6001,-34.6355 23.1988,-34.2521 21.8083,-33.8241 20.4306,-33.3528 19.0681,-32.8394 17.7233,-32.2853 16.3989,-31.692 15.0979,-31.0612 13.8236,-30.3946 12.5794,-29.6941 11.3691,-28.9615 10.1967,-28.1988 9.0662,-27.4082 7.982,-26.5918 6.9486,-25.7519 5.9706,-24.8907 5.0526,-24.0106 4.1995,-23.114 3.4159,-22.2035 2.7067,-21.2814 2.0765,-20.3503 1.53,-19.4127 1.0715,-18.4713 .7053,-17.5287 .4355,-16.5873 .2657,-15.6497 .1993,-14.7186 .2395,-13.7965 .3888,-12.886 .6495,-11.9894 1.0232,-11.1093 1.5112,-10.2481 2.1141,-9.4082 2.8321,-8.5918 3.6647,-7.8012 4.6109,-7.0385 5.6689,-6.3059 6.8365,-5.6054 8.1108,-4.9388 9.4884,-4.308 10.9651,-3.7147 12.5362,-3.1606 14.1966,-2.6472 15.9404,-2.1759 17.7614,-1.7479 19.6528,-1.3645 21.6074,-1.0266 23.6176,-.7353 25.6755,-.4913 27.7728,-.2953 29.901,-.1479 32.0515,-.0493 34.2154,0 36.3837,0 38.5476,-.0493\"/>\n",
       "<text text-anchor=\"middle\" x=\"35.2995\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">start</text>\n",
       "</g>\n",
       "<!-- pre&#45;process -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>pre&#45;process</title>\n",
       "<ellipse fill=\"none\" stroke=\"#000000\" cx=\"172.2432\" cy=\"-18\" rx=\"65.7887\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"172.2432\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\" fill=\"#000000\">pre&#45;process</text>\n",
       "</g>\n",
       "<!-- _start&#45;&gt;pre&#45;process -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>_start&#45;&gt;pre&#45;process</title>\n",
       "<path fill=\"none\" stroke=\"#000000\" d=\"M69.8562,-18C78.0608,-18 87.1345,-18 96.3902,-18\"/>\n",
       "<polygon fill=\"#000000\" stroke=\"#000000\" points=\"96.4726,-21.5001 106.4726,-18 96.4726,-14.5001 96.4726,-21.5001\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.dot.Digraph at 0x7f2d73584c90>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "project = mlrun.get_or_create_project(\"myproj\", \"./\")\n",
    "\n",
    "fn_echo = project.set_function(\"echo_function\", kind=\"serving\", image=\"mlrun/mlrun\")\n",
    "\n",
    "graph_echo = fn_echo.set_topology(\"flow\")\n",
    "\n",
    "graph_echo.to(class_name=\"Echo\", name=\"pre-process\", some_arg=\"abc\")\n",
    "\n",
    "graph_echo.plot(rankdir=\"LR\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a mock server to test this locally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': '97397ea412334afdb5e4cb7d7c2e6dd3'}\n",
      "Echo: pre-process {'inputs': 123}\n"
     ]
    }
   ],
   "source": [
    "echo_server = fn_echo.to_mock_server(current_function=\"*\")\n",
    "\n",
    "result = echo_server.test(\"\", {\"inputs\": 123})\n",
    "echo_server.wait_for_completion()\n",
    "\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**For more information, see the [Advanced model serving graph notebook example](./graph-example.ipynb).**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<!-- Requires better description - and title? -->\n",
    "\n",
    "You can use any Python function by specifying the handler name (e.g. `handler=json.dumps`).\n",
    "The function is triggered with the `event.body` as the first argument, and its result \n",
    "is passed to the next step.\n",
    "\n",
    "Alternatively, you can use classes that can also store some step/configuration and separate the \n",
    "one time init logic from the per event logic. The classes are initialized with the `class_args`.\n",
    "If the class init args contain `context` or `name`, they are initialized with the \n",
    "[graph context](./realtime-pipelines.ipynb) and the step name. \n",
    "\n",
    "By default, the `class_name` and handler specify a class/function name in the `globals()` (i.e. this module).\n",
    "Alternatively, those can be full paths to the class (module.submodule.class), e.g. `storey.WriteToParquet`.\n",
    "You can also pass the module as an argument to functions such as `function.to_mock_server(namespace=module)`.\n",
    "In this case the class or handler names are also searched in the provided module.\n",
    " <!-- the previous sentence needs clarification -->\n",
    "\n",
    "When using classes the class event handler is invoked on every event with the `event.body`.\n",
    "If the Task step `full_event` parameter is set to `True` the handler is invoked and returns\n",
    "the full `event` object. If the class event handler is not specified, it invokes the class `do()` method. \n",
    "\n",
    "If you need to implement async behavior, then subclass `storey.MapClass`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Create a single step\n",
    "\n",
    "When creating a serving function with a step that is a part of the class, the graph context is created in the init, and you can use this context to get the functions' parameters.\n",
    "\n",
    "When creating a single step (no class), the context is not created, and therefore the `get_param` does not work. The following example illustrates how to create the context and then use the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting serving-handler-func.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile serving-handler-func.py\n",
    "import pandas as pd\n",
    "import mlrun\n",
    "import mlrun.utils\n",
    "import os\n",
    "import json\n",
    "\n",
    "def test(event):\n",
    "    serving_spec = mlrun.utils.get_serving_spec()\n",
    "    server_context = mlrun.serving.GraphServer().from_dict(serving_spec)\n",
    "    context = mlrun.serving.GraphContext(server=server_context)\n",
    "    param = context.get_param(\"Test\")\n",
    "    return param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "serving_func_handler = project.set_function(\n",
    "    name=\"serving-handler-func\",\n",
    "    func=\"serving-handler-func.py\",\n",
    "    image=\"mlrun/mlrun\",\n",
    "    kind=\"serving\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "serving_func_handler.spec.parameters = {\"Test\": \"test\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = serving_func_handler.set_topology(\"flow\")\n",
    "\n",
    "graph.to(name=\"test\", handler=\"test\").respond()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> 2023-05-09 14:24:55,287 [info] Starting remote function deploy\n",
      "2023-05-09 14:24:56  (info) Deploying function\n",
      "2023-05-09 14:24:56  (info) Building\n",
      "2023-05-09 14:24:56  (info) Staging files and preparing base images\n",
      "2023-05-09 14:24:57  (info) Building processor image\n",
      "2023-05-09 14:26:02  (info) Build complete\n",
      "> 2023-05-09 14:26:08,232 [info] successfully deployed function: {'internal_invocation_urls': ['nuclio-serving-context-shapira-serving-handler-func.default-tenant.svc.cluster.local:8080'], 'external_invocation_urls': ['serving-context-shapira-serving-handler-func-serving-c-6v7nqbg6.default-tenant.app.cust-cs-il-3-5-2.iguazio-cd2.com/']}\n"
     ]
    }
   ],
   "source": [
    "serving_func_deploy = project.deploy_function(\"serving-handler-func\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "serving_func_deploy.function.invoke(\n",
    "    \"/\",\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
